#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>

void flip_bit(uint8_t *buf, size_t byte_idx, int bit) {
    buf[byte_idx] ^= (1u << bit);
}

unsigned long matmul_sum(int N, int do_inject) {
    uint64_t *A = calloc(N*N, sizeof(uint64_t));
    uint64_t *B = calloc(N*N, sizeof(uint64_t));
    uint64_t *C = calloc(N*N, sizeof(uint64_t));
    if (!A || !B || !C) return 0;

    for (int i=0;i<N*N;i++) {
        A[i] = i + 1;
        B[i] = i + 2;
        C[i] = 0;
    }

    for (int i=0;i<N;i++)
        for (int j=0;j<N;j++) {
        uint64_t sum = 0;
        for (int k=0;k<N;k++)
            sum += A[i*N + k] * B[k*N + j];
        C[i*N + j] = sum;
    }

    char *inj = getenv("FI_INJECT");
    if (inj && atoi(inj) == 1 && do_inject) {
        char *off_s = getenv("FI_OFFSET");
        char *bit_s = getenv("FI_BIT");
        if (off_s && bit_s) {
            size_t byte_off = (size_t) atol(off_s);
            int bit = atoi(bit_s);
            uint8_t *raw = (uint8_t*) C;
            size_t total_bytes = N*N*sizeof(uint64_t);
            if (byte_off < total_bytes && bit >= 0 && bit < 8) {
                flip_bit(raw, byte_off, bit);
            }
        }
    }

    unsigned long sum = 0;
    for (int i=0;i<N*N;i++)
        sum += C[i];
    free(A); free(B); free(C);
    return sum;
}

int main(int argc, char **argv) {
    char *mode = getenv("FI_MODE");
    int N = 64;
    if (mode && strcmp(mode, "defense") == 0) {
        unsigned long r1 = matmul_sum(N, 1);
        unsigned long r2 = matmul_sum(N, 1);
        if (r1 != r2) {
            printf("DETECTED_MISMATCH\n");
            return 2;
        } else {
            printf("%lu\n", r1);
            return 0;
        }
    } else {
        unsigned long r = matmul_sum(N, 1);
        printf("%lu\n", r);
        return 0;
    }
}
